import os
import cv2
import numpy as np
from PIL import Image
from typing import List, Dict, Tuple, Optional

 
class FaceRecognitionSystem:
    def __init__(self, face_cascade_path: str = "haarcascade_frontalface_default.xml"):
        """初始化人脸识别系统"""
        # 加载人脸检测器
        self.face_cascade = cv2.CascadeClassifier(face_cascade_path)
        # 加载人脸识别器
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        # 人脸数据库：ID到名称的映射
        self.id_to_name = {}

    def _preprocess_image(self, image_path: str) -> np.ndarray:
        """预处理图像：转换为灰度图并检测人脸"""
        # 读取图像
        image = cv2.imread(image_path)
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 检测人脸
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
        if len(faces) == 0:
            raise ValueError("未在图像中检测到人脸")
        # 假设图像中只有一个人脸，返回第一个检测到的人脸
        (x, y, w, h) = faces[0]
        return gray[y:y + h, x:x + w]

    def register_face(self, image_path: str, person_name: str, person_id: int) -> None:
        """注册人脸到人脸库"""
        try:
            # 预处理图像
            face_roi = self._preprocess_image(image_path)
            # 保存人脸ID和名称的映射
            self.id_to_name[person_id] = person_name
            # 训练识别器（实际应用中可能需要收集多幅图像）
            self.recognizer.update([face_roi], np.array([person_id]))
            print(f"已成功注册人脸：{person_name} (ID: {person_id})")
        except Exception as e:
            print(f"注册人脸失败：{str(e)}")

    def recognize_face(self, image_path: str, threshold: float = 100.0) -> Tuple[Optional[str], float]:
        """识别图像中的人脸"""
        try:
            # 预处理图像
            face_roi = self._preprocess_image(image_path)
            # 预测人脸
            person_id, confidence = self.recognizer.predict(face_roi)
            # 计算相似度（置信度越低，相似度越高）
            similarity = max(0, 100 - confidence)

            if similarity >= threshold:
                person_name = self.id_to_name.get(person_id, f"未知(ID:{person_id})")
                return person_name, similarity
            else:
                return None, similarity  # 未识别出人脸
        except Exception as e:
            print(f"人脸识别失败：{str(e)}")
            return None, 0.0

    def detect_faces_in_image(self, image_path: str, output_path: Optional[str] = None) -> List[
        Tuple[int, int, int, int]]:
        """检测图像中的所有人脸并返回位置"""
        # 读取图像
        image = cv2.imread(image_path)
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 检测人脸
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        # 如果指定了输出路径，则绘制人脸框并保存图像
        if output_path:
            for (x, y, w, h) in faces:
                cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 0), 2)
            cv2.imwrite(output_path, image)

        return faces

    def save_model(self, model_path: str) -> None:
        """保存训练好的模型"""
        self.recognizer.save(model_path)
        print(f"模型已保存至：{model_path}")

    def load_model(self, model_path: str, id_mapping: Dict[int, str]) -> None:
        """加载预训练模型"""
        self.recognizer.read(model_path)
        self.id_to_name = id_mapping
        print(f"已加载模型：{model_path}")


# 示例用法
if __name__ == "__main__":
    # 创建人脸识别系统实例
    fr_system = FaceRecognitionSystem()

    # 注册人脸
    fr_system.register_face("path/to/person1.jpg", "张三", 1)
    fr_system.register_face("path/to/person2.jpg", "李四", 2)

    # 保存模型
    fr_system.save_model("face_model.yml")

    # 检测图像中的人脸
    faces = fr_system.detect_faces_in_image("test_image.jpg", "output.jpg")
    print(f"在图像中检测到 {len(faces)} 个人脸")

    # 识别人脸
    person_name, similarity = fr_system.recognize_face("test_image.jpg", threshold=70.0)
    if person_name:
        print(f"识别结果：{person_name}，相似度：{similarity:.2f}%")
    else:
        print("未识别出人脸")